package cn.jbolt.core.Interceptor.reqlimit.redis;

import cn.jbolt.core.Interceptor.reqlimit.JBoltRateResult;
import cn.jbolt.core.base.config.JBoltConfig;
import cn.jbolt.core.cache.JBoltCacheType;
import com.jfinal.kit.StrKit;
import com.jfinal.plugin.redis.Redis;

import java.util.Arrays;

public class JBoltRedisRateLimitUtil {
    /**
     * 限流lua脚本
     */
    private static final String luaScript =
            "local lockKey = KEYS[1]\n" +
            "local lockCount = ARGV[1]\n" +
            "local lockExpire = ARGV[2]\n" +
            "local currentCount = tonumber(redis.call('get',lockKey))\n" +
            "if currentCount == nil or currentCount < tonumber(lockCount) then\n" +
            "redis.call(\"INCRBY\", lockKey,\"1\")\n" +
            "redis.call(\"expire\",lockKey,lockExpire)\n" +
            "return 1\n" +
            "else\n" +
            "return 0\n" +
            "end";

    public static JBoltRateResult limit(String key, int time,int count,String msg) {
        if(StrKit.isBlank(luaScriptSha)){
            initRedisLuaScriptSha();
        }
        return Redis.use(JBoltConfig.JBOLT_CACHE_NAME).call(jedis -> {
            Object evalRet = jedis.evalsha(luaScriptSha,
                                            Arrays.asList(key),
                                            Arrays.asList(String.valueOf(count),String.valueOf(time)));
            if(evalRet == null|| evalRet.toString().equals("0")){
                return JBoltRateResult.notPass(StrKit.defaultIfBlank(msg,JBoltRateResult.DEFAULT_MSG));
            }
            return JBoltRateResult.PASS;
        });

    }

    private static String luaScriptSha = null;
    private static String luaScriptShaRedisKey = "jbolt_redis_lua_rate_limit_script";
    /**
     * 系统启动加载redis lua脚本
     */
    public static void initRedisLuaScriptSha() {
        if(!JBoltCacheType.REDIS.equals(JBoltConfig.JBOLT_CACHE_TYPE)) {
            return;
        }
        //从缓存里拿到指定的sha
        luaScriptSha = Redis.use(JBoltConfig.JBOLT_CACHE_NAME).call(jedis -> {
            //获取sha
            String sha = jedis.get(luaScriptShaRedisKey);
            //如果拿到了 得去判断脚本在不在
            if(StrKit.notBlank(sha)) {
                //判断sha指向的脚本是否存在
                boolean exists = jedis.scriptExists(luaScriptSha);
                if (exists) {
                    //有sha 有脚本存在 就返回sha
                    return sha;
                }
            }
            //如果没有 就初始化脚本
            return jedis.scriptLoad(luaScript);
        });

        if(StrKit.isBlank(luaScriptSha)){
            throw new RuntimeException("启动时加载redis+lua限流脚本失败");
        }
    }
}
